# 1. 导入依赖库
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import (roc_auc_score, f1_score, precision_score, recall_score,
                             roc_curve, confusion_matrix, accuracy_score)
from xgboost import XGBClassifier
from imblearn.over_sampling import SMOTE
import warnings

warnings.filterwarnings("ignore")  # 忽略无关警告

# 设置中文字体（避免画图乱码）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 1. 数据读取与基础探索（仅探索Train数据，Final Test不参与探索）
def load_and_explore_data(train_path="../data/train.csv", final_test_path="../data/test.csv"):
    """
    读取数据：
    - train.csv：用于划分训练/验证集，做基础探索
    - test.csv：最终测试集，仅读取不探索（留到最后预测）
    """
    # 读取训练数据（用于拆分和探索）
    train_df = pd.read_csv(train_path)
    # 读取最终测试数据（留到最后，暂不处理目标列）
    final_test_df = pd.read_csv(final_test_path)

    print(f"=== 数据基本信息 ===")
    print(f"训练数据（train.csv）形状: {train_df.shape}")
    print(f"最终测试数据（test.csv）形状: {final_test_df.shape}")

    # 仅探索训练数据的缺失值和目标分布
    print(f"\n=== 训练数据缺失值统计 ===")
    print(f"训练集缺失值总数: {train_df.isnull().sum().sum()}")
    print(f"最终测试集缺失值总数: {final_test_df.isnull().sum().sum()}")  # 仅检查缺失，不处理

    # 无效特征识别（全量数据统一删除：员工编号、全员18+、标准工时）
    drop_cols = ["EmployeeNumber", "Over18", "StandardHours"]
    train_df = train_df.drop(drop_cols, axis=1)
    # 最终测试集保留员工编号（用于最后输出结果），仅删除Over18和StandardHours
    final_test_df = final_test_df.drop(["Over18", "StandardHours"], axis=1)

    # 训练数据目标变量分布（流失率）
    print(f"\n=== 训练数据目标变量（Attrition）分布 ===")
    attrition_count = train_df["Attrition"].value_counts()
    print(attrition_count)
    print(f"训练数据流失率: {attrition_count[1] / len(train_df):.2%}")

    # 可视化训练数据目标分布（仅画训练集）
    plt.figure(figsize=(6, 4))
    sns.countplot(x="Attrition", data=train_df)
    plt.title("训练数据员工流失分布（0=留存，1=流失）")
    plt.show()

    return train_df, final_test_df


# 2. 特征工程（训练集和最终测试集统一应用，避免数据泄露）
def feature_engineering(df, is_final_test=False):
    """
    构造业务新特征：
    - is_final_test：标记是否为最终测试集（若为True，不处理Attrition列）
    """
    df_fe = df.copy()
    # 1. 收入效率：月收入/总工作年限（+1避免除以0）
    df_fe["IncomePerYear"] = df_fe["MonthlyIncome"] / (df_fe["TotalWorkingYears"] + 1)
    # 2. 通勤成本：离家距离/工作时长
    df_fe["DistanceRatio"] = df_fe["DistanceFromHome"] / (df_fe["YearsAtCompany"] + 1)
    # 3. 晋升速度：当前职位在职时长/上次升职时间
    df_fe["PromotionGap"] = df_fe["YearsInCurrentRole"] / (df_fe["YearsSinceLastPromotion"] + 1)
    # 4. 综合满意度：工作+环境+关系满意度的平均值
    df_fe["SatisfactionScore"] = (df_fe["JobSatisfaction"] + df_fe["EnvironmentSatisfaction"] +
                                  df_fe["RelationshipSatisfaction"]) / 3

    return df_fe


# 3. 数据预处理（异常值+编码，仅用训练集拟合，避免泄露）
def preprocess_data(train_df, final_test_df):
    """
    预处理流程：
    1. 拆分训练集为「训练子集（X_train, y_train）」和「验证子集（X_val, y_val）」
    2. 异常值处理（仅数值特征，IQR方法）
    3. 特征编码（数值标准化+分类独热编码，仅用训练子集拟合）
    4. 最终测试集仅做transform，不拟合
    """
    # 3.1 拆分训练集与验证集（8:2，stratify保持流失率一致）
    X_train_full = train_df.drop("Attrition", axis=1)
    y_train_full = train_df["Attrition"]
    X_train_split, X_val_split, y_train_split, y_val_split = train_test_split(
        X_train_full, y_train_full, test_size=0.2, random_state=42, stratify=y_train_full
    )
    print(f"\n=== 训练/验证集拆分结果 ===")
    print(f"训练子集样本数: {len(X_train_split)} (流失率: {y_train_split.mean():.2%})")
    print(f"验证子集样本数: {len(X_val_split)} (流失率: {y_val_split.mean():.2%})")

    # 3.2 定义数值特征和分类特征（基于训练集）
    num_features = ["Age", "DistanceFromHome", "MonthlyIncome", "TotalWorkingYears",
                    "YearsAtCompany", "YearsInCurrentRole", "YearsSinceLastPromotion",
                    "YearsWithCurrManager", "NumCompaniesWorked", "TrainingTimesLastYear",
                    "IncomePerYear", "DistanceRatio", "PromotionGap", "SatisfactionScore"]  # 新增特征加入
    # 不是数值类的特征都是分类特征
    cat_features = [col for col in X_train_split.columns if col not in num_features]

    # 3.3 异常值处理（仅对数值特征，IQR截断，避免极端值影响）
    def handle_outliers(df, num_features):
        df_clean = df.copy()
        for col in num_features:
            Q1 = df_clean[col].quantile(0.25)
            Q3 = df_clean[col].quantile(0.75)
            IQR = Q3 - Q1
            print(f"{col}: IQR={IQR:.2f}, Q1={Q1:.2f}, Q3={Q3:.2f}")
            # 收入和总工作年限允许更大的波动
            if col in ["MonthlyIncome", "TotalWorkingYears"]:
                upper_bound = Q3 + 3 * IQR  # 收入允许更大波动
            else:
                upper_bound = Q3 + 1.5 * IQR
            lower_bound = Q1 - 1.5 * IQR
            df_clean[col] = df_clean[col].clip(lower=lower_bound, upper=upper_bound)
        return df_clean

    # 处理训练子集和验证子集的异常值（最终测试集后续处理）
    X_train_split_clean = handle_outliers(X_train_split, num_features)
    X_val_split_clean = handle_outliers(X_val_split, num_features)

    # 3.4 特征编码流水线（仅用训练子集拟合！）
    # 数值特征：标准化
    scaler = StandardScaler()
    # 分类特征：onehot编码（drop="first"避免共线性，忽略未知类别）
    onehot = OneHotEncoder(drop="first", sparse_output=False, handle_unknown="ignore")

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", scaler, num_features),
            ("cat", onehot, cat_features)
        ],
        remainder="passthrough"
    )

    # 训练子集：fit+transform；验证子集/最终测试集：仅transform（避免泄露）
    X_train_prep = preprocessor.fit_transform(X_train_split_clean)
    X_val_prep = preprocessor.transform(X_val_split_clean)

    # 处理最终测试集（提取特征，保留员工编号用于输出）
    final_test_employee_id = final_test_df["EmployeeNumber"].copy()  # 保留员工编号
    X_final_test = final_test_df.drop(["EmployeeNumber", "Attrition"], axis=1, errors="ignore")  # 移除ID和目标列
    X_final_test_clean = handle_outliers(X_final_test, num_features)  # 异常值处理
    X_final_test_prep = preprocessor.transform(X_final_test_clean)  # 仅transform

    # 获取编码后的特征名
    cat_feature_names = preprocessor.named_transformers_["cat"].get_feature_names_out(cat_features)
    all_feature_names = num_features + list(cat_feature_names)

    return (X_train_prep, y_train_split, X_val_prep, y_val_split,
            X_final_test_prep, final_test_employee_id, preprocessor, all_feature_names)


# 4. 处理数据不平衡（仅对训练子集过采样，验证集保持原始分布）
def handle_imbalance(X_train_prep, y_train_split):
    """
    SMOTE过采样：仅对训练子集（X_train_split）处理，验证集不处理
    目的：避免模型偏向多数类（留存员工），同时保持验证集真实性
    """
    smote = SMOTE(random_state=42, k_neighbors=3)
    X_train_smote, y_train_smote = smote.fit_resample(X_train_prep, y_train_split)

    print(f"\n=== SMOTE过采样结果 ===")
    print(f"过采样前训练子集：留存{sum(y_train_split == 0)}个，流失{sum(y_train_split == 1)}个")
    print(f"过采样后训练子集：留存{sum(y_train_smote == 0)}个，流失{sum(y_train_smote == 1)}个")

    return X_train_smote, y_train_smote


# 5. 模型训练与超参数调优（XGBoost，基于训练/验证集）
def train_and_optimize_model(X_train_smote, y_train_smote, X_val_prep, y_val_split):
    """
    模型训练流程：
    1. XGBoost初始模型
    2. 网格搜索调优超参数（5折交叉验证）
    3. 验证集评估模型性能
    4. 可视化ROC曲线
    """
    # 初始模型（二分类任务，评估指标AUC）
    xgb = XGBClassifier(
        objective="binary:logistic",  # 二分类任务(流失/不流失),输出概率
        eval_metric="auc",  # AUC作为评估指标
        random_state=42,  # 随机数种子
        n_jobs=-1  # 多线程加速训练
    )

    # 超参数网格（平衡调优效果与速度）
    # param_grid = {
    #     "learning_rate": [0.01, 0.05, 0.1],  # 学习率调优:控制每棵树的贡献权重,小值更精确但需要更多树
    #     "n_estimators": [100, 200, 300],  # 数的棵树数量调优:控制模型复杂度，防止过拟合
    #     "max_depth": [3, 5, 7],  # 数的深度调优:控制树的最大深度，防止过拟合
    #     "subsample": [0.7, 0.8, 0.9],  # 行采样:每棵树随机选部分样本训练(避免过拟合)
    #     "colsample_bytree": [0.7, 0.8, 0.9],  # 列采样:每棵树随机选部分特征进行训练(避免过拟合)
    #     "min_child_weight": [1, 3, 5],  # 控制叶节点最小样本权重和（防止过拟合）
    # }
    param_grid = {
        "learning_rate": [0.05],  # 学习率调优:控制每棵树的贡献权重,小值更精确但需要更多树
        "n_estimators": [200],  # 数的棵树数量调优:控制模型复杂度，防止过拟合
        "max_depth": [7],  # 数的深度调优:控制树的最大深度，防止过拟合
        "subsample": [0.8],  # 行采样:每棵树随机选部分样本训练(避免过拟合)
        "colsample_bytree": [0.7],  # 列采样:每棵树随机选部分特征进行训练(避免过拟合)
        "min_child_weight": [1],  # 控制叶节点最小样本权重和（防止过拟合）
    }

    # 网格搜索（用训练子集过采样后的数据，验证集评估）
    grid_search = GridSearchCV(
        estimator=xgb,  # 模型
        param_grid=param_grid,  # 超参数组合
        cv=5,  # 5折交叉验证(把数据集分成5份，轮流当验证集,评估更稳定)
        scoring="roc_auc",  # 用AUC作为评分标准(越大模型效果越好)
        n_jobs=-1,  # 多线程加速搜索
        verbose=1  # 输出搜索过程
    )
    grid_search.fit(X_train_smote, y_train_smote)  # 用过采样后的训练子集训练模型

    # 最优模型
    best_xgb = grid_search.best_estimator_
    print(f"\n=== 最优模型参数 ===")
    print(grid_search.best_params_)
    # 验证集预测与评估
    y_val_pred = best_xgb.predict(X_val_prep)
    y_val_pred_proba = best_xgb.predict_proba(X_val_prep)[:, 1]

    print(f"\n验证集模型评估")
    print(f"Accuracy（准确率）: {accuracy_score(y_val_split, y_val_pred):.4f}")
    print(f"AUC: {roc_auc_score(y_val_split, y_val_pred_proba):.4f}")
    print(f"Recall（召回率）: {recall_score(y_val_split, y_val_pred):.4f}")
    print(f"Precision（精确率）: {precision_score(y_val_split, y_val_pred):.4f}")
    print(f"F1 Score: {f1_score(y_val_split, y_val_pred):.4f}")

    # 可视化验证集ROC曲线
    fpr, tpr, _ = roc_curve(y_val_split, y_val_pred_proba)
    plt.figure(figsize=(6, 4))
    plt.plot(fpr, tpr, color="blue", label=f"AUC = {roc_auc_score(y_val_split, y_val_pred_proba):.4f}")
    plt.plot([0, 1], [0, 1], "k--", label="随机猜测")
    plt.xlabel("假正例率（FPR）")
    plt.ylabel("真正例率（TPR）")
    plt.title("验证集ROC曲线")
    plt.legend()
    plt.savefig("val_roc_curve.png", dpi=300, bbox_inches="tight")
    plt.close()

    return best_xgb

# 6. 最终测试集预测（用最优模型，输出最终结果）
def predict_final_test(best_xgb, X_final_test_prep, final_test_employee_id, final_test_df, thresholds=0.29):
    """
    对最终测试集做预测，增加准确率（Accuracy）计算
    """
    # 最终预测概率与标签（基于传入的阈值）
    final_pred_proba = best_xgb.predict_proba(X_final_test_prep)[:, 1].round(4)
    final_pred = np.where(final_pred_proba >= thresholds, 1, 0)  # 基于阈值的预测标签

    # 风险等级划分
    def get_risk_level(prob):
        if prob >= 0.3:
            return "高风险（需优先干预）"
        else:
            return "低风险（无需特殊处理）"

    # 构造最终结果表
    final_result = pd.DataFrame({
        "EmployeeNumber": final_test_employee_id,
        "Attrition_Pred_Probability": final_pred_proba,
        "Attrition_Prediction": final_pred,
        "Risk_Level": [get_risk_level(p) for p in final_pred_proba]
    })

    # 若有真实标签，计算并输出准确率（Accuracy）及其他指标
    if "Attrition" in final_test_df.columns:
        final_result["Attrition_True"] = final_test_df["Attrition"].values
        # 计算核心指标
        test_accuracy = accuracy_score(final_test_df["Attrition"], final_pred)  # 准确率
        test_f1 = f1_score(final_test_df["Attrition"], final_pred)  # F1
        test_recall = recall_score(final_test_df["Attrition"], final_pred)  # 召回率
        test_precision = precision_score(final_test_df["Attrition"], final_pred)
        test_auc = roc_auc_score(final_test_df["Attrition"], final_pred)

        # 输出所有指标
        print(f"\n最终测试集评估:")
        print(f"准确率（Accuracy）: {test_accuracy:.4f}")  # 新增：整体预测正确的比例
        print(f"召回率（Recall）: {test_recall:.4f}")
        print(f"精确率（Precision）: {test_precision:.4f}")
        print(f"F1 Score: {test_f1:.4f}")
        print(f"AUC: {test_auc:.4f}")
        print(f"预测流失人数: {sum(final_pred)}, 真实流失人数: {sum(final_test_df['Attrition'])}")
    else:
        print(f"最终测试集无真实标签，无法计算指标")
        print(f"预测流失人数: {sum(final_pred)}, 预测留存人数: {len(final_pred) - sum(final_pred)}")

    # 保存结果
    final_result.to_csv("final_test_attrition_prediction.csv", index=False, encoding="utf-8-sig")

    return final_result


# 9. 主函数（串联全流程，确保最终测试集留到最后）
def main():
    # 1.读取数据（训练集+最终测试集）
    train_df, final_test_df = load_and_explore_data(
        train_path="../data/train.csv",
        final_test_path="../data/test.csv"
    )

    # 2.特征工程,构造可能会有影响的新的特征（训练集+最终测试集统一处理）
    train_df_fe = feature_engineering(train_df, is_final_test=False)
    final_test_df_fe = feature_engineering(final_test_df, is_final_test=True)

    # 3.数据预处理（拆分训练/验证集，仅用训练集拟合）
    (X_train_prep, y_train_split, X_val_prep, y_val_split,
     X_final_test_prep, final_test_employee_id, preprocessor, feature_names) = preprocess_data(
        train_df_fe, final_test_df_fe
    )

    # 4.处理训练子集不平衡（SMOTE过采样）
    X_train_smote, y_train_smote = handle_imbalance(X_train_prep, y_train_split)

    # 5.模型训练与调优（基于训练/验证集）
    best_xgb = train_and_optimize_model(X_train_smote, y_train_smote, X_val_prep, y_val_split)

    # 6.最终测试集预测（留到最后，输出结果）
    final_result = predict_final_test(best_xgb, X_final_test_prep, final_test_employee_id, final_test_df)


# 运行主函数
if __name__ == "__main__":
    main()
